//! 参见[`Data2D`]、[`Array2D`]、[`Slice2D`]、[`SliceMut2D`]
//!
//! TODO: 等到自定义 DST 能用了就换成它

use self::{
    builder::Builder,
    succ::{CoordReducer, IndexReducer, LineReducer, ZipReducer},
};
use crate::{
    geometry::{Coord, Dir},
    repr::{BoxRepr, MutRepr, RefRepr, ReprMutUsz, ReprRefUsz, ReprUsz},
    serde_utils::IterSeq,
};
#[cfg(test)]
use proptest::prelude::*;
use serde::{
    de::{DeserializeSeed, Error, Visitor},
    Deserialize, Deserializer, Serialize, Serializer,
};
use std::{
    cell::Cell,
    fmt::{self, Debug},
    hash::{Hash, Hasher},
    marker::PhantomData,
    num::NonZeroUsize,
    ops::{Index, IndexMut},
};

/// 非空的二维数组
pub struct Data2D<'a, T: 'a, R: ReprUsz> {
    cols: NonZeroUsize,
    /// 长度必然为列数乘以行数，非空
    ///
    /// ```text
    ///    |- cols=4 -|
    ///  +--0--1--2--3-- x
    /// 0|  0  1  2  3
    /// 1|  4  5  6  7
    /// 2|  8  9 10 11
    ///  y
    /// ```
    data: R::WrapUsz<'a, [T]>,
}

/// 其生命周期为泛型，而非 `'static`。
///
/// 这是因为 GAT 一定不变，不能是协变或逆变。
/// 于是，要允许存储不是 `'static` 的数据，必须使用泛型。
///
/// 缺点是，在类型上，泛型必须一层层地传递上去。
/// 实际使用时，包括 `impl`，通常可以省略。
pub type Array2D<'a, T> = Data2D<'a, T, BoxRepr>;
pub type Slice2D<'a, T> = Data2D<'a, T, RefRepr>;
pub type SliceMut2D<'a, T> = Data2D<'a, T, MutRepr>;

impl<T> Array2D<'_, T> {
    /// 如果是空的就返回空
    pub fn new(size: Coord, val: T) -> Option<Self>
    where
        T: Clone,
    {
        (size.y != 0).then_some(())?;
        Some(Data2D {
            cols: NonZeroUsize::new(size.x)?,
            data: vec![val; size.x * size.y].into_boxed_slice(),
        })
    }

    /// 如果是空的就返回空。会克隆一遍。
    pub fn from_iter<'a>(iter: impl IntoIterator<Item = &'a [T]>) -> Option<Self>
    where
        T: Clone + 'a,
    {
        let mut b = Builder::new();
        for row in iter {
            b.push_many(row);
            b.end_line().then_some(())?;
        }
        b.build()
    }

    /// 任意变化生命周期
    ///
    /// TODO：也许可以搞一个 `ReprOwn` trait
    pub fn vary<'a>(self) -> Array2D<'a, T> {
        Data2D {
            cols: self.cols,
            data: self.data,
        }
    }

    pub fn deserialize_with<'a, U, D>(
        mut map: impl FnMut(U) -> T,
        deserializer: D,
    ) -> Result<Self, D::Error>
    where
        U: Deserialize<'a>,
        D: Deserializer<'a>,
    {
        struct Vis<'a, T, U, F, const IS_ROWS: bool> {
            builder: &'a mut Builder<T>,
            map: &'a mut F,
            _marker: PhantomData<U>,
        }

        impl<'a, 'de, T, U, F, const IS_ROWS: bool> Visitor<'de> for Vis<'a, T, U, F, IS_ROWS>
        where
            F: FnMut(U) -> T,
            U: Deserialize<'de>,
        {
            type Value = ();

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                formatter.write_str(if IS_ROWS { "rows" } else { "columns" })
            }

            fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
            where
                A: serde::de::SeqAccess<'de>,
            {
                if IS_ROWS {
                    while let Some(()) = seq.next_element_seed(Vis::<_, _, _, false> {
                        builder: self.builder,
                        map: self.map,
                        _marker: PhantomData,
                    })? {}
                } else {
                    while let Some(el) = seq.next_element()? {
                        self.builder.push((self.map)(el));
                    }
                    if !self.builder.end_line() {
                        return Err(A::Error::custom("格式错误"));
                    }
                }
                Ok(())
            }
        }

        impl<'a, 'de, T, U, F, const IS_ROWS: bool> DeserializeSeed<'de> for Vis<'a, T, U, F, IS_ROWS>
        where
            F: FnMut(U) -> T,
            U: Deserialize<'de>,
        {
            type Value = ();

            fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
            where
                D: Deserializer<'de>,
            {
                deserializer.deserialize_seq(self)
            }
        }

        let mut builder = Builder::new();
        Vis::<_, _, _, true> {
            builder: &mut builder,
            map: &mut map,
            _marker: PhantomData,
        }
        .deserialize(deserializer)?;
        builder.build().ok_or_else(|| D::Error::custom("格式错误"))
    }
}

impl<T, R: ReprUsz> Data2D<'_, T, R> {
    /// 即使返回 `Some`，也可能不存在对应下标
    fn get_index(&self, Coord { x, y }: Coord) -> Option<usize> {
        let cols = self.cols.get();
        (x < cols).then_some(x + y * cols)
    }
}

impl<T, R: ReprRefUsz> Data2D<'_, T, R> {
    pub fn size(&self) -> Coord {
        let cols = self.cols.get();
        Coord::new(cols, self.raw_data().len() / cols)
    }

    pub fn get(&self, index: Coord) -> Option<&T> {
        let index = self.get_index(index)?;
        self.raw_data().get(index)
    }

    pub fn borrow(&self) -> Slice2D<T> {
        Data2D {
            cols: self.cols,
            data: self.raw_data(),
        }
    }

    fn rows(&self) -> impl Iterator<Item = &[T]> + Clone {
        self.raw_data().chunks(self.cols.get())
    }

    /// 顺序不指定
    pub fn raw_data(&self) -> &[T] {
        R::borrow_unsized(&self.data)
    }

    fn lines<L: LineReducer>(
        &self,
        dir: Dir,
    ) -> impl Iterator<Item = impl Iterator<Item = L::T> + Clone> + Clone {
        let size = self.size();
        let step = move |dir: Dir| move |prev: &_| L::step(size, dir, prev);
        // 每行的开头
        let head = {
            let order = dir.axis().rotate().face(dir.positive());
            let rows = size[order.axis()];
            let begin = if dir.positive() {
                L::init(size)
            } else {
                L::init_rev(size)
            };
            std::iter::successors(Some(begin), step(order)).take(rows)
        };
        let cols = size[dir.axis()];
        head.map(move |head| std::iter::successors(Some(head), step(dir)).take(cols))
    }

    pub fn walk_lines(
        &self,
        dir: Dir,
    ) -> impl Iterator<Item = impl Iterator<Item = Coord> + Clone> + Clone {
        self.lines::<CoordReducer>(dir)
    }

    pub fn iter_lines(
        &self,
        dir: Dir,
    ) -> impl Iterator<Item = impl Iterator<Item = &T> + Clone + '_> + Clone + '_ {
        let data = self.raw_data();
        self.lines::<IndexReducer>(dir)
            .map(move |line| line.map(|idx| &data[idx]))
    }

    pub fn enum_lines(
        &self,
        dir: Dir,
    ) -> impl Iterator<Item = impl Iterator<Item = (Coord, &T)> + Clone + '_> + Clone + '_ {
        let data = self.raw_data();
        self.lines::<ZipReducer<CoordReducer, IndexReducer>>(dir)
            .map(|line| line.map(|(c, idx)| (c, &data[idx])))
    }

    pub fn serialize_with<'a, F, U, S>(&'a self, serializer: S, map: F) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
        U: Serialize,
        F: Fn(&'a T) -> U,
    {
        serializer.collect_seq(self.rows().map(|row| IterSeq(row.iter().map(&map))))
    }
}

impl<T, R: ReprMutUsz> Data2D<'_, T, R> {
    /// 顺序不指定
    pub fn raw_data_mut(&mut self) -> &mut [T] {
        R::borrow_mut_unsized(&mut self.data)
    }

    pub fn get_mut(&mut self, index: Coord) -> Option<&mut T> {
        let index = self.get_index(index)?;
        self.raw_data_mut().get_mut(index)
    }

    pub fn as_cells(&mut self) -> Slice2D<Cell<T>> {
        Data2D {
            cols: self.cols,
            data: Cell::from_mut(self.raw_data_mut()).as_slice_of_cells(),
        }
    }

    pub fn borrow_mut(&mut self) -> SliceMut2D<T> {
        Data2D {
            cols: self.cols,
            data: self.raw_data_mut(),
        }
    }
}

impl<'a, T, R: ReprUsz> Clone for Data2D<'a, T, R>
where
    R::WrapUsz<'a, [T]>: Clone,
{
    fn clone(&self) -> Self {
        Data2D {
            cols: self.cols,
            data: self.data.clone(),
        }
    }
}

impl<T, R: ReprRefUsz> Index<Coord> for Data2D<'_, T, R> {
    type Output = T;
    fn index(&self, index: Coord) -> &Self::Output {
        self.get(index).expect("下标不存在")
    }
}
impl<T, R: ReprMutUsz> IndexMut<Coord> for Data2D<'_, T, R> {
    fn index_mut(&mut self, index: Coord) -> &mut Self::Output {
        self.get_mut(index).expect("下标不存在")
    }
}

impl<T: Debug, R: ReprRefUsz> Debug for Data2D<'_, T, R> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_list().entries(self.rows()).finish()
    }
}

impl<T: Hash, R: ReprRefUsz> Hash for Data2D<'_, T, R> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.cols.hash(state);
        self.raw_data().hash(state);
    }
}

impl<T: PartialEq, R: ReprRefUsz> PartialEq for Data2D<'_, T, R> {
    fn eq(&self, other: &Self) -> bool {
        self.cols == other.cols && self.raw_data() == other.raw_data()
    }
}
impl<T: Eq, R: ReprRefUsz> Eq for Data2D<'_, T, R> {}

impl<T: Serialize, R: ReprRefUsz> Serialize for Data2D<'_, T, R> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        self.serialize_with(serializer, |el| el)
    }
}

impl<'de, T: Deserialize<'de>> Deserialize<'de> for Array2D<'_, T> {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        Self::deserialize_with(|e| e, deserializer)
    }
}

mod builder {
    use super::Array2D;
    use std::num::NonZeroUsize;

    pub struct Builder<T> {
        cols: Option<NonZeroUsize>,
        data: Vec<T>,
        row_stop: usize,
    }

    impl<T> Builder<T> {
        pub fn new() -> Self {
            Self {
                cols: None,
                data: Vec::new(),
                row_stop: 0,
            }
        }

        pub fn push(&mut self, val: T) {
            self.data.push(val);
        }

        /// 会克隆
        pub fn push_many(&mut self, slice: &[T])
        where
            T: Clone,
        {
            self.data.extend_from_slice(slice);
        }

        /// 如果返回 `false`，则什么都没有发生。不过不建议继续使用。
        pub fn end_line(&mut self) -> bool {
            match self.cols {
                Some(cols) => (self.data.len() == cols.get() + self.row_stop)
                    .then(|| self.row_stop = self.data.len())
                    .is_some(),
                None => NonZeroUsize::new(self.data.len())
                    .map(|len| {
                        self.cols = Some(len);
                        self.row_stop = len.get();
                    })
                    .is_some(),
            }
        }

        pub fn build<'a>(self) -> Option<Array2D<'a, T>> {
            (self.data.len() == self.row_stop && self.row_stop != 0).then_some(())?;
            Some(Array2D {
                cols: self.cols?,
                data: self.data.into_boxed_slice(),
            })
        }
    }
}

mod succ {
    use crate::geometry::{Coord, Dir};
    use std::marker::PhantomData;

    pub trait LineReducer {
        type T: Clone;
        fn init(size: Coord) -> Self::T;
        fn init_rev(size: Coord) -> Self::T;
        fn step(size: Coord, dir: Dir, prev: &Self::T) -> Option<Self::T>;
    }

    pub struct CoordReducer {}
    impl LineReducer for CoordReducer {
        type T = Coord;
        fn init(_: Coord) -> Self::T {
            Coord::new(0, 0)
        }
        fn init_rev(size: Coord) -> Self::T {
            size - Coord::new(1, 1)
        }
        fn step(_: Coord, dir: Dir, &prev: &Self::T) -> Option<Self::T> {
            prev + dir
        }
    }

    pub struct IndexReducer {}
    impl LineReducer for IndexReducer {
        type T = usize;
        fn init(_: Coord) -> Self::T {
            0
        }
        fn init_rev(Coord { x, y }: Coord) -> Self::T {
            x * y - 1
        }
        fn step(size: Coord, dir: Dir, &prev: &Self::T) -> Option<Self::T> {
            let cols = size.x as isize;
            let offset = match dir {
                Dir::Down => cols,
                Dir::Left => -1,
                Dir::Right => 1,
                Dir::Up => -cols,
            };
            prev.checked_add_signed(offset)
        }
    }

    pub struct ZipReducer<T1: LineReducer, T2: LineReducer> {
        _marker: PhantomData<(T1, T2)>,
    }
    impl<T1: LineReducer, T2: LineReducer> LineReducer for ZipReducer<T1, T2> {
        type T = (T1::T, T2::T);
        fn init(size: Coord) -> Self::T {
            (T1::init(size), T2::init(size))
        }
        fn init_rev(size: Coord) -> Self::T {
            (T1::init_rev(size), T2::init_rev(size))
        }
        fn step(size: Coord, dir: Dir, prev: &Self::T) -> Option<Self::T> {
            T1::step(size, dir, &prev.0).zip(T2::step(size, dir, &prev.1))
        }
    }
}

#[cfg(test)]
impl<T: Debug + Arbitrary> Arbitrary for Array2D<'_, T> {
    type Parameters = ();
    type Strategy = BoxedStrategy<Self>;
    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
        ((1..15usize), (1..15usize))
            .prop_map(|(cols, rows)| (cols, cols * rows))
            .prop_flat_map(|(cols, size)| {
                (Just(cols), prop::collection::vec(any::<T>(), size..=size))
            })
            .prop_map(|(cols, data)| Self {
                cols: cols.try_into().unwrap(),
                data: data.into_boxed_slice(),
            })
            .boxed()
    }
}

#[cfg(test)]
mod test {
    use super::Array2D;
    use crate::geometry::{Coord, Dir};
    use itertools::Itertools;
    use proptest::{prop_assert, prop_assert_eq, proptest};
    use serde_json::{from_str, to_string};
    use std::fmt::Debug;

    #[test]
    fn test() {
        let mut arr = Array2D::new(Coord::new(3, 2), 0).unwrap();
        arr[(0, 0).into()] = 1;
        arr[(1, 0).into()] = 4;
        arr[(2, 0).into()] = 5;
        arr.borrow_mut()[(0, 1).into()] = 9;
        arr.borrow_mut()[(1, 1).into()] = 8;

        fn asteq<T: PartialEq + Debug>(
            i: impl IntoIterator<Item = impl IntoIterator<Item = T>>,
            c: impl IntoIterator<Item = impl IntoIterator<Item = T>>,
        ) {
            for (i, c) in i.into_iter().zip(c.into_iter()) {
                for (i, c) in i.into_iter().zip(c.into_iter()) {
                    debug_assert_eq!(i, c);
                }
            }
        }

        asteq(arr.iter_lines(Dir::Left), &[[0, 8, 9], [5, 4, 1]]);
        asteq(arr.iter_lines(Dir::Down), &[[1, 9], [4, 8], [5, 0]]);
        asteq(
            arr.borrow().walk_lines(Dir::Right),
            [
                [(0, 0).into(), (1, 0).into(), (2, 0).into()],
                [(0, 1).into(), (1, 1).into(), (2, 1).into()],
            ],
        );
        asteq(
            arr.borrow_mut().enum_lines(Dir::Up),
            [
                [((2, 1).into(), &0), ((2, 0).into(), &5)],
                [((1, 1).into(), &8), ((1, 0).into(), &4)],
                [((0, 1).into(), &9), ((0, 0).into(), &1)],
            ],
        );
    }

    #[test]
    fn test_serde() {
        assert!(Array2D::<()>::from_iter([]).is_none());
        assert!(Array2D::from_iter([&[] as &[()]]).is_none());
        assert!(Array2D::from_iter([&[1, 9, 1, 9] as &[_], &[8, 1, 0]]).is_none());

        let arr = Array2D::from_iter([1, 1, 4, 5, 1, 4].chunks(3)).unwrap();
        assert_eq!(to_string(&arr).unwrap(), "[[1,1,4],[5,1,4]]");
        assert_eq!(arr, from_str("[[1,1,4],[5,1,4]]").unwrap());
    }

    proptest!(
        #[test]
        fn fuzz(arr: Array2D<u8>) {
            assert_eq!(arr, from_str(&to_string(&arr).unwrap()).unwrap());

            for dir in Dir::EACH {
                prop_assert!(arr.walk_lines(dir).flatten().all_unique());

                for (a @ (c, &v), b) in arr
                    .walk_lines(dir)
                    .flatten()
                    .zip(arr.iter_lines(dir).flatten())
                    .zip(arr.enum_lines(dir).flatten())
                {
                    prop_assert_eq!(a, b);
                    prop_assert_eq!(arr[c], v);
                }
            }
        }
    );
}
